#!/home/muyi/Tool/.venv/bin/python3

import asyncio
import subprocess
import ipaddress
import socket
from gglobal import config,logger

# 尝试导入netifaces，如果失败则使用备用方法
try:
    import netifaces
    HAS_NETIFACES = True
except ImportError:
    HAS_NETIFACES = False
    logger.debug("Warning: netifaces not available, 将使用备用网段检测方法")

# 尝试导入asyncssh，如果失败则使用简化版本
try:
    import asyncssh
    HAS_ASYNCSSH = True
except ImportError:
    HAS_ASYNCSSH = False
    logger.debug("Warning: asyncssh not available, SSH功能将被禁用")

# 限制并发任务的数量 - 提高并发数以加快扫描速度
MAX_CONCURRENT_TASKS = 260
# 移除全局semaphore，改为在函数内部创建

# 优化超时设置
PING_TIMEOUT_MS = 2000  # ping超时时间(毫秒)
PORT_TIMEOUT_SEC = 2  # 端口检测超时时间(秒)
SSH_TIMEOUT_SEC = 5    # SSH连接超时时间(秒)


def get_local_ip_networks():
    """获取本机所有网络接口的IP网段
    
    返回:
        list: 本机IP网段列表，如 ['192.168.1', '10.0.0']
    """
    networks = []
    
    if HAS_NETIFACES:
        try:
            # 方法1: 使用netifaces获取所有网络接口
            for interface in netifaces.interfaces():
                try:
                    addrs = netifaces.ifaddresses(interface)
                    if netifaces.AF_INET in addrs:
                        for addr_info in addrs[netifaces.AF_INET]:
                            ip = addr_info.get('addr')
                            netmask = addr_info.get('netmask')
                            if ip and netmask and not ip.startswith('127.'):
                                # 计算网段
                                network = ipaddress.IPv4Network(f"{ip}/{netmask}", strict=False)
                                network_prefix = '.'.join(str(network.network_address).split('.')[:-1])
                                if network_prefix not in networks:
                                    networks.append(network_prefix)
                except Exception as e:
                    continue
        except Exception:
            pass
    
    # 方法2: 备用方法，使用socket连接外部地址获取本机IP
    if not networks:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.connect(("8.8.8.8", 80))
                local_ip = s.getsockname()[0]
                network_prefix = '.'.join(local_ip.split('.')[:-1])
                networks.append(network_prefix)
        except Exception:
            pass
    
    # 方法3: 使用socket.gethostbyname获取本机IP
    if not networks:
        try:
            hostname = socket.gethostname()
            local_ip = socket.gethostbyname(hostname)
            if not local_ip.startswith('127.'):
                network_prefix = '.'.join(local_ip.split('.')[:-1])
                networks.append(network_prefix)
        except Exception:
            pass
    
    # 方法4: Windows系统使用ipconfig命令
    if not networks:
        try:
            import platform
            if platform.system().lower() == 'windows':
                # Windows系统：添加CREATE_NO_WINDOW标志隐藏命令行窗口
                result = subprocess.run(
                    ['ipconfig'], 
                    capture_output=True, 
                    text=True, 
                    timeout=5,
                    creationflags=subprocess.CREATE_NO_WINDOW
                )
                lines = result.stdout.split('\n')
                for line in lines:
                    if 'IPv4' in line and ':' in line:
                        ip = line.split(':')[-1].strip()
                        if ip and not ip.startswith('127.') and '.' in ip:
                            try:
                                # 验证IP格式
                                ipaddress.IPv4Address(ip)
                                network_prefix = '.'.join(ip.split('.')[:-1])
                                if network_prefix not in networks:
                                    networks.append(network_prefix)
                            except:
                                continue
        except Exception:
            pass
    
    # 如果没有找到任何网段，使用默认值
    if not networks:
        networks = ['192.168.1', '192.168.0', '10.0.0']
        
    return networks


def get_primary_local_network():
    """获取主要的本机IP网段
    
    返回:
        str: 主要的IP网段，如 '192.168.1'
    """
    networks = get_local_ip_networks()
    
    # 优先返回192.168.x网段
    for network in networks:
        if network.startswith('192.168.'):
            return network
            
    # 其次返回10.x.x网段
    for network in networks:
        if network.startswith('10.'):
            return network
            
    # 返回第一个找到的网段
    return networks[0] if networks else '192.168.1'


async def ping_scan(ip, semaphore):
    async with semaphore:
        # Windows和Linux的ping命令参数不同
        import platform
        if platform.system().lower() == 'windows':
            # Windows系统：使用更兼容的方式隐藏窗口
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            startupinfo.wShowWindow = subprocess.SW_HIDE
            
            process = await asyncio.create_subprocess_exec(
                'ping', '-n', '1', '-w', str(PING_TIMEOUT_MS), ip,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                startupinfo=startupinfo
            )
        else:
            process = await asyncio.create_subprocess_exec(
                'ping', '-c', '1', '-W', '1', ip,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
        
        try:
            # 添加整体超时控制
            stdout, stderr = await asyncio.wait_for(
                process.communicate(), 
                timeout=PING_TIMEOUT_MS / 1000 + 0.5
            )
            is_success = process.returncode == 0
            
            # 记录ping结果日志
            if is_success:
                # logger.debug(f"Ping {ip} 成功 - 返回码: {process.returncode}")
                # 记录ping输出的关键信息（避免日志过长）
                if stdout:
                    stdout_str = stdout.decode('utf-8', errors='ignore').strip()
                    # 只记录包含时间信息的行
                    for line in stdout_str.split('\n'):
                        if 'time=' in line.lower() or '时间=' in line:
                            logger.debug(f"Ping {ip} 响应: {line.strip()}")
                            break
            else:
                # logger.debug(f"Ping {ip} 失败 - 返回码: {process.returncode}")
                if stderr:
                    stderr_str = stderr.decode('utf-8', errors='ignore').strip()
                    if stderr_str:
                        logger.debug(f"Ping {ip} 错误信息: {stderr_str[:100]}...")  # 限制错误信息长度
            
            return is_success
        except asyncio.TimeoutError:
            # 超时则终止进程
            # logger.debug(f"Ping {ip} 超时 - 超过 {PING_TIMEOUT_MS}ms")
            try:
                process.terminate()
                await process.wait()
            except:
                pass
            return False


async def check_port(ip, port=None, semaphore=None):
    async with semaphore:
        # 从配置获取端口参数，如果没有传入参数的话
        if port is None:
            port = config.get_param('port', 22)
            
        try:
            # 添加连接超时控制
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(ip, port),
                timeout=PORT_TIMEOUT_SEC
            )
            writer.close()
            await writer.wait_closed()
            return True
        except (asyncio.TimeoutError, ConnectionRefusedError, OSError):
            return False
        except Exception:
            return False


async def get_robot_id(ip, username=None, password=None, semaphore=None):
    async with semaphore:
        if not HAS_ASYNCSSH:
            return None
        
        # 从配置获取SSH参数，如果没有传入参数的话
        if username is None:
            username = config.get_param('username', 'muyi')
        if password is None:
            password = config.get_param('password', 'muyi@123')
        port = config.get_param('port', 22)
            
        try:
            # 添加SSH连接超时控制
            conn = await asyncio.wait_for(
                asyncssh.connect(
                    ip, 
                    port=port, 
                    username=username, 
                    password=password,
                    known_hosts=None,  # 跳过主机密钥验证以加快连接
                    client_keys=None   # 不使用密钥认证以加快连接
                ),
                timeout=SSH_TIMEOUT_SEC
            )
            
            async with conn:
                # 添加命令执行超时控制
                result = await asyncio.wait_for(
                    conn.run('cat /mooe/machineinfo/robot_id', check=True),
                    timeout=SSH_TIMEOUT_SEC
                )
                # 处理文件内容
                cleaned_data = '\n'.join(line.strip() for line in result.stdout.splitlines() if line.strip())
                return cleaned_data
                
        except (asyncio.TimeoutError, ConnectionRefusedError, OSError):
            return None
        except Exception:
            # 对于SSH相关的异常，直接返回None而不打印错误
            return None


async def scan_ip_range(start_ip, end_ip, ssh_user, ssh_pass):
    # 创建当前事件循环的semaphore
    semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)
    
    start_ip = ipaddress.ip_address(start_ip)
    end_ip = ipaddress.ip_address(end_ip)
    ip_range = ipaddress.summarize_address_range(start_ip, end_ip)

    active_hosts = []
    hosts_with_robot_id = []
    tasks = []

    try:
        for network in ip_range:
            for ip in network:
                ip_str = str(ip)
                ping_task = ping_scan(ip_str, semaphore)
                port_task = check_port(ip_str, None, semaphore)  # 使用配置中的端口
                tasks.append((ip_str, asyncio.gather(ping_task, port_task)))

        results = await asyncio.gather(*[task[1] for task in tasks], return_exceptions=True)

        for (ip_str, result) in zip((task[0] for task in tasks), results):
            if isinstance(result, Exception):
                continue
            is_active, port_open = result
            if is_active:
                if port_open:
                    robot_id = await get_robot_id(ip_str, ssh_user, ssh_pass, semaphore)
                    if robot_id:
                        active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID: {robot_id}")
                        hosts_with_robot_id.append((robot_id, ip_str))
                    else:
                        active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID not found")
                else:
                    active_hosts.append(f"IP: {ip_str} - Port 22 closed")
    except Exception as e:
        logger.debug(f"扫描过程中发生错误: {e}")

    return active_hosts, hosts_with_robot_id


async def scan_local_network(network_range=None, ssh_user="muyi", ssh_pass="muyi@123", fast_mode=True):
    """扫描局域网IP段
    
    参数:
        network_range: 网络段，支持 "192.168.1" 或 "192.168.1.0/24" 格式。如果为None，则自动检测本机网段
        ssh_user: SSH用户名
        ssh_pass: SSH密码
        fast_mode: 快速模式，启用后会使用批量处理和早期结果返回来提高速度
    
    返回:
        tuple: (active_hosts, hosts_with_robot_id)
    """
    # 创建当前事件循环的semaphore
    semaphore = asyncio.Semaphore(MAX_CONCURRENT_TASKS)
    
    # 如果没有指定网络段，自动检测本机网段
    if network_range is None:
        network_range = get_primary_local_network()
        logger.debug(f"自动检测到本机网段: {network_range}")
    
    active_hosts = []
    hosts_with_robot_id = []
    tasks = []

    try:
        # 解析网络段格式
        ip_list = []
        
        if '/' in network_range:
            # CIDR格式，如 "10.10.17.0/24"
            try:
                network = ipaddress.IPv4Network(network_range, strict=False)
                # 跳过网络地址和广播地址
                ip_list = [str(ip) for ip in network.hosts()]
                logger.debug(f"CIDR格式网段 {network_range} 包含 {len(ip_list)} 个主机地址")
            except ValueError as e:
                logger.debug(f"无效的CIDR格式: {network_range}, 错误: {e}")
                return active_hosts, hosts_with_robot_id
        else:
            # 传统格式，如 "192.168.1"
            for i in range(1, 255):
                ip_list.append(f"{network_range}.{i}")
            logger.debug(f"传统格式网段 {network_range} 包含 {len(ip_list)} 个主机地址")
        
        # 智能扫描策略：优先扫描常用IP段以提高发现设备的速度
        def get_priority_ips(ip_list):
            """按优先级排序IP地址，常用IP段优先"""
            priority_ips = []
            normal_ips = []
            
            for ip in ip_list:
                last_octet = int(ip.split('.')[-1])
                # 优先扫描常用的IP段：1-20, 100-120, 200-220
                if (1 <= last_octet <= 20) or (100 <= last_octet <= 120) or (200 <= last_octet <= 220):
                    priority_ips.append(ip)
                else:
                    normal_ips.append(ip)
            
            return priority_ips + normal_ips
        
        # 按优先级重新排序IP列表
        if not ('/' in network_range):  # 只对传统格式进行优化排序
            ip_list = get_priority_ips(ip_list)
        
        # 快速模式：分批处理以提高响应性和内存效率
        if fast_mode:
            batch_size = 50  # 每批处理50个IP
            for i in range(0, len(ip_list), batch_size):
                batch_ips = ip_list[i:i + batch_size]
                batch_tasks = []
                
                # 创建当前批次的任务
                for ip_str in batch_ips:
                    ping_task = ping_scan(ip_str, semaphore)
                    port_task = check_port(ip_str, None, semaphore)  # 使用配置中的端口
                    batch_tasks.append((ip_str, asyncio.gather(ping_task, port_task)))
                
                # 执行当前批次
                batch_results = await asyncio.gather(*[task[1] for task in batch_tasks], return_exceptions=True)
                
                # 处理当前批次的结果
                for (ip_str, result) in zip((task[0] for task in batch_tasks), batch_results):
                    if isinstance(result, Exception):
                        continue
                    is_active, port_open = result
                    if is_active:
                        if port_open:
                            robot_id = await get_robot_id(ip_str, ssh_user, ssh_pass, semaphore)
                            if robot_id:
                                active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID: {robot_id}")
                                hosts_with_robot_id.append((robot_id, ip_str))
                            else:
                                active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID not found")
                        else:
                            active_hosts.append(f"IP: {ip_str} - Port 22 closed")
                
                # 输出当前进度
                progress = min(i + batch_size, len(ip_list))
                logger.debug(f"扫描进度: {progress}/{len(ip_list)} ({progress/len(ip_list)*100:.1f}%)")
        else:
            # 传统模式：一次性处理所有IP
            for ip_str in ip_list:
                ping_task = ping_scan(ip_str, semaphore)
                port_task = check_port(ip_str, None, semaphore)  # 使用配置中的端口
                tasks.append((ip_str, asyncio.gather(ping_task, port_task)))

            results = await asyncio.gather(*[task[1] for task in tasks], return_exceptions=True)

            for (ip_str, result) in zip((task[0] for task in tasks), results):
                if isinstance(result, Exception):
                    continue
                is_active, port_open = result
                if is_active:
                    if port_open:
                        robot_id = await get_robot_id(ip_str, ssh_user, ssh_pass, semaphore)
                        if robot_id:
                            active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID: {robot_id}")
                            hosts_with_robot_id.append((robot_id, ip_str))
                        else:
                            active_hosts.append(f"IP: {ip_str} - Port 22 open, Robot ID not found")
                    else:
                        active_hosts.append(f"IP: {ip_str} - Port 22 closed")
    except Exception as e:
        logger.debug(f"扫描过程中发生错误: {e}")

    return active_hosts, hosts_with_robot_id
